import numpy as np
import pandas as pd
import scanpy as sc

from scipy.sparse import issparse
from scipy.sparse.linalg import ArpackNoConvergence
from scipy.sparse.linalg import svds


# 01 生成模糊矩阵（Fuzzy Matrix）和行权重（Dc）
def _mca_step1(X: np.ndarray) -> dict:
	"""
	mca_step 1: 对基因表达矩阵进行预处理，生成模糊矩阵（Fuzzy Matrix）和行权重（Dc）。
	参数:
		X: 基因表达矩阵 (genes x cells), numpy 数组或稀疏矩阵。
	返回:
		dict: 包含 Z(模糊矩阵)和 Dc(行权重)的字典。
	"""
	# 如果 X 是稀疏矩阵，则转换为稠密矩阵。
	if issparse(X):
		AM = X.toarray()
	else:
		AM = X
	
	rmin = np.min(AM, axis=1, keepdims=True)
	rmax = np.max(AM, axis=1, keepdims=True)
	range_vals = rmax - rmin
	
	AM -= rmin
	AM = np.divide(AM, np.where(range_vals != 0, range_vals, 1.0), dtype=np.float64)
	
	FM = np.vstack([AM, 1 - AM])
	
	total = np.sum(FM)
	colsum = np.sum(FM, axis=0)
	rowsum = np.sum(FM, axis=1)
	colsum_safe = np.where(colsum == 0, 1.0, colsum)
	rowsum_safe = np.where(rowsum == 0, 1.0, rowsum)
	
	FM /= np.sqrt(colsum_safe)
	FM /= np.sqrt(rowsum_safe)[:, np.newaxis]
	Dc = 1 / np.sqrt(rowsum_safe / total)
	
	dict_Z_Dc = {
			"Z":  FM,
			"Dc": Dc
	}
	
	return dict_Z_Dc


def _mca_step2(
		Z: np.ndarray,
		V: np.ndarray,
		Dc: np.ndarray
) -> dict:
	"""
	mca_step 2：计算细胞坐标和基因坐标。
	参数:
		Z: 模糊矩阵 (from mca_step1, shape: 2n_genes x n_cells)
		V: SVD 的右奇异向量 (n_cells x components)
		Dc: 行权重 (from mca_step1, length: n_genes)
	返回:
		dict: 包含细胞坐标和基因坐标的字典
	"""
	AZ = Z
	AV = V
	ADc = Dc.reshape(-1, 1)  # 2n_genes x 1
	
	features_coordinates = AZ @ AV
	
	features_coordinates *= ADc
	n_genes = features_coordinates.shape[0] // 2
	features_coordinates = features_coordinates[:n_genes, :]
	
	AZcol = AZ.shape[1]
	cells_coordinates = np.sqrt(AZcol) * AV
	
	dict_cells_features = {
			"cellsCoordinates":    cells_coordinates,
			"featuresCoordinates": features_coordinates
	}
	
	return dict_cells_features


def run_mca(
		adata: sc.AnnData,
		nmcs: int = 64,
		features: list | None = None,
		meta: bool = False,
) -> tuple[np.ndarray, np.ndarray, np.ndarray] | tuple[pd.DataFrame, pd.DataFrame, np.ndarray]:
	"""
	run_mca: 整合 mca 分析流程，从 AnnData 对象中提取数据并返回降维结果。
	参数:
		adata: AnnData 对象，包含基因表达矩阵 X
		nmcs: 降维后的维度
		features: 需要分析的基因列表，如果为 None，则选择所有基因
		include_meta: 是否包含 meta 信息，如果为 True，则返回 DataFrame 形式的细胞坐标和基因载荷信息
	返回:
		tuple: 包含细胞坐标、基因载荷、标准差的元组，或者包含细胞坐标、基因载荷 DataFrame 以及标准差的元组
	"""
	print("run mca...")
	
	if features is not None:
		adata = adata[:, adata.var_names & features]
	
	X = adata.X
	if issparse(X):
		X = X.toarray()
	X = X.T
	
	print("mca step 1: Construct the Fuzzy Matrix and Row Weights")
	step1_result = _mca_step1(X)
	print("mca step 1 completed: Fuzzy Matrix and Row Weights constructed")
	
	print("svd started")
	num_components = nmcs + 1
	try:
		u, s, vt = svds(step1_result["Z"], k=num_components, which='LM')
		if len(s) < num_components:
			raise ValueError(f"Only {len(s)} singular values converged, requested {num_components}")
	except ArpackNoConvergence as e:
		print(f"ARPACK did not converge: {e}")
		if hasattr(e, 'eigenvalues'):
			u, s, vt = e.eigenvectors, e.eigenvalues, e.eigenvectors.T
			print(f"Using partial results with {len(s)} components")
		else:
			raise
	except ValueError as e:
		print(f"Value error in svds: {e}")
		raise
	except Exception as e:
		print(f"Unexpected error in svds: {e}")
		raise
	
	sort_indices = np.argsort(s)[::-1]
	s_desc = s[sort_indices]
	v_desc = vt.T[:, sort_indices]
	V = v_desc[:, 1:]
	print("svd completed")
	
	print("mca step 2: Calculate Cell Coordinates and Feature Loadings")
	step2_result = _mca_step2(step1_result["Z"], V, step1_result["Dc"])
	cell_embedding = step2_result["cellsCoordinates"]
	gene_loading = step2_result["featuresCoordinates"]
	stdev = s_desc[1:]
	print("mca step 2 completed: Cell Coordinates and Feature Loadings calculated")
	
	if meta:
		cell_embedding_df = pd.DataFrame(cell_embedding, index=adata.obs_names)
		cell_embedding_df.columns = [f"MC_{col}" for col in cell_embedding_df.columns]
		gene_loading_df = pd.DataFrame(gene_loading, index=adata.var_names)
		gene_loading_df.index.name = "Gene"
		gene_loading_df.columns = [f"MC_{col}" for col in gene_loading_df.columns]
		
		return cell_embedding_df, gene_loading_df, stdev
	
	return cell_embedding, gene_loading, stdev
